#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Token生成器 - 基于UUID和雪花算法
"""

import time
import uuid
import threading
from datetime import datetime, timedelta
from jose import jwt, JWTError
from typing import Optional, Dict, Any
import config
from src.utils.logger import get_logger
logger = get_logger("capcut")

class SnowflakeGenerator:
    """雪花算法ID生成器"""
    
    def __init__(self, datacenter_id: int = 1, worker_id: int = 1):
        # 基准时间戳 (2024-01-01 00:00:00)
        self.epoch = 1704067200000  # 毫秒
        
        # 各部分位数
        self.datacenter_id_bits = 5
        self.worker_id_bits = 5
        self.sequence_bits = 12
        
        # 最大值
        self.max_datacenter_id = (1 << self.datacenter_id_bits) - 1
        self.max_worker_id = (1 << self.worker_id_bits) - 1
        self.max_sequence = (1 << self.sequence_bits) - 1
        
        # 位移
        self.worker_id_shift = self.sequence_bits
        self.datacenter_id_shift = self.sequence_bits + self.worker_id_bits
        self.timestamp_shift = (self.sequence_bits + 
                               self.worker_id_bits + 
                               self.datacenter_id_bits)
        
        # 验证参数
        if datacenter_id > self.max_datacenter_id or datacenter_id < 0:
            raise ValueError(f"datacenter_id 必须在 0-{self.max_datacenter_id} 之间")
        if worker_id > self.max_worker_id or worker_id < 0:
            raise ValueError(f"worker_id 必须在 0-{self.max_worker_id} 之间")
        
        self.datacenter_id = datacenter_id
        self.worker_id = worker_id
        self.sequence = 0
        self.last_timestamp = -1
        
        # 线程锁
        self._lock = threading.Lock()
    
    def _get_timestamp(self) -> int:
        """获取当前时间戳（毫秒）"""
        return int(time.time() * 1000)
    
    def _wait_next_millis(self, last_timestamp: int) -> int:
        """等待下一毫秒"""
        timestamp = self._get_timestamp()
        while timestamp <= last_timestamp:
            timestamp = self._get_timestamp()
        return timestamp
    
    def next_id(self) -> int:
        """生成下一个ID"""
        with self._lock:
            timestamp = self._get_timestamp()
            
            # 时钟回退检查
            if timestamp < self.last_timestamp:
                raise RuntimeError(f"时钟回退，拒绝生成ID。当前时间戳: {timestamp}, 上次时间戳: {self.last_timestamp}")
            
            # 同一毫秒内
            if timestamp == self.last_timestamp:
                self.sequence = (self.sequence + 1) & self.max_sequence
                if self.sequence == 0:
                    # 序列号用完，等待下一毫秒
                    timestamp = self._wait_next_millis(self.last_timestamp)
            else:
                self.sequence = 0
            
            self.last_timestamp = timestamp
            
            # 生成ID
            snowflake_id = (
                ((timestamp - self.epoch) << self.timestamp_shift) |
                (self.datacenter_id << self.datacenter_id_shift) |
                (self.worker_id << self.worker_id_shift) |
                self.sequence
            )
            
            return snowflake_id


class TokenGenerator:
    """Token生成器"""
    
    def __init__(self):
        self.snowflake = SnowflakeGenerator(datacenter_id=1, worker_id=1)
        self.secret_key = config.JWT_SECRET_KEY
        self.algorithm = config.JWT_ALGORITHM
        self.expire_minutes = config.JWT_ACCESS_TOKEN_EXPIRE_MINUTES
    
    def generate_token_id(self) -> str:
        """生成简短的Token ID"""
        # 生成短UUID (8位)
        short_uuid = str(uuid.uuid4()).replace('-', '')[:8]
        
        # 生成雪花ID并转换为16进制 (取后8位)
        snowflake_id = self.snowflake.next_id()
        short_snowflake = format(snowflake_id & 0xFFFFFFFF, '08x')
        
        # 组合生成最终的token_id (16位)
        token_id = f"{short_uuid}{short_snowflake}"
        return token_id
    
    def create_access_token(self, data: Dict[str, Any], expires_delta: Optional[timedelta] = None) -> str:
        """创建访问令牌"""
        to_encode = data.copy()
        
        # 设置过期时间
        if expires_delta:
            expire = datetime.utcnow() + expires_delta
        else:
            expire = datetime.utcnow() + timedelta(minutes=self.expire_minutes)
        
        # 生成唯一Token ID
        token_id = str(data.get('user_id'))

        logger.info(f"生成Token ID: {token_id}")
        
        # 添加标准JWT声明
        to_encode.update({
            "exp": expire,
            "iat": datetime.utcnow(),
            "jti": token_id,  # JWT ID
            "iss": "capcut-api",  # 签发者
            "aud": "capcut-client"  # 受众
        })
        
        # 生成JWT
        encoded_jwt = jwt.encode(to_encode, self.secret_key, algorithm=self.algorithm)
        logger.info(f"生成Token: {encoded_jwt}")
        return encoded_jwt
    
    def verify_token(self, token: str) -> Optional[Dict[str, Any]]:
        """验证令牌"""
        try:
            payload = jwt.decode(
                token, 
                self.secret_key, 
                algorithms=[self.algorithm],
                audience="capcut-client",
                issuer="capcut-api"
            )
            logger.info(f"验证Token成功: {payload}")
            return payload
        except JWTError as e:
            logger.error(f"验证Token失败: {e}",exc_info=True)
            return None
    
    def decode_token_info(self, token: str) -> Optional[Dict[str, Any]]:
        """解码Token信息（不验证签名，用于调试）"""
        try:
            # 不验证签名，仅解码
            payload = jwt.decode(
                token, 
                options={"verify_signature": False}
            )
            return payload
        except Exception:
            return None


# 全局Token生成器实例
token_generator = TokenGenerator() 